#!/usr/bin/python2.4
# -*- coding: utf-8 -*-
#
# Copyright 2006, 2007 Google Inc. All Rights Reserved.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2, or (at your option)
# any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

"""Unittest for yawt module"""

import StringIO
import os
import tempfile
import unittest
import yawt

class YARTestCase(unittest.TestCase):
  """Extends unittest.TestCase with two commonly used functions:
    OpenFileObject() and VerifyListAndRead()"""

  def OpenFileObject(self, content, mode=None, binary_encoding=None):
    """Create a YAR object using a file-like object

    Args:
      content: the content of the file-like object
      mode: a mode to pass to YAR, mode can be: 'r', 'w', 'a'
      binary_encoding: the default encoding to use for binary files

    Returns:
      A YAR object
    """
    fileobj = StringIO.StringIO(content)
    if mode and binary_encoding:
      return yawt.YAR(fileobj, mode=mode, binary_encoding=binary_encoding)
    elif mode:
      return yawt.YAR(fileobj, mode=mode)
    if binary_encoding:
      return yawt.YAR(fileobj, binary_encoding=binary_encoding)
    else:
      return yawt.YAR(fileobj)

  def VerifyListAndRead(self, yar_obj, file_list, file_contents_dict):
    """Verify the list of files from the YAR file, and then check to
    make sure those files can be read correctly.

    Args:
      yar_obj: a YAR object to check
      file_list: the exact list of (file, type) tuples yar_obj should contain
      file_contents_dict: the dictionary with the file name to file content
        mapping to check
    """
    # read the file data
    yar_obj.yar_file.seek(0)
    yar_data = yar_obj.yar_file.read()
    # create a new yar file
    new_yar_obj = yawt.YAR(StringIO.StringIO(yar_data))
    self.assertEqual(new_yar_obj.List(), file_list)
    for (f, t) in file_list:
      self.assertEqual(new_yar_obj.Read(f), file_contents_dict[f])


class YARStringObjectTestCase(YARTestCase):
  def setUp(self):
    self.file_list = [('test.txt', 'file')]
    self.file_contents = {'test.txt': 'hello world\n'}
    (fd, self.tmpname) = tempfile.mkstemp('yartest')
    os.close(fd)

  def tearDown(self):
    os.remove(self.tmpname)

  def WriteTempFile(self, content):
    """Write content to a temporary file

    Args:
      content: the content of the temporary file
    """
    tmpfile = open(self.tmpname, 'wb')
    tmpfile.write(content)
    tmpfile.close()

  def OpenFileString(self, content, mode=None, binary_encoding=None):
    """Create a YAR object using a filename, with a temporary file

    Args:
      content: the content of the temporary file
      mode: a mode to pass to YAR, mode can be: 'r', 'w', 'a'
      binary_encoding: the default encoding to use for binary files

    Returns:
      A YAR object
    """
    self.WriteTempFile(content)
    if mode and binary_encoding:
      return yawt.YAR(self.tmpname, mode=mode, binary_encoding=binary_encoding)
    elif mode:
      return yawt.YAR(self.tmpname, mode=mode)
    if binary_encoding:
      return yawt.YAR(self.tmpname, binary_encoding=binary_encoding)
    else:
      return yawt.YAR(self.tmpname)


class YARConstructorOpenTest(YARStringObjectTestCase):
  def testOpenYarFile(self):
    yar_content = ('<YAR 1>\n'
        '<file name="test.txt">\n'
        'hello world\n'
        '</file>\n')
    yar_string = self.OpenFileString(yar_content)
    yar_object = self.OpenFileObject(yar_content)
    for yar in yar_string, yar_object:
      self.VerifyListAndRead(yar, self.file_list, self.file_contents)
      yar.Close()

  def testOpenDosYarFile(self):
    yar_content = ('<YAR 1>\r\n'
        '<file name="test.txt">\r\n'
        'hello world\r\n'
        '</file>\r\n')
    yar_string = self.OpenFileString(yar_content)
    yar_object = self.OpenFileObject(yar_content)
    for yar in yar_string, yar_object:
      self.VerifyListAndRead(yar, self.file_list, self.file_contents)
      yar.Close()

  def testInvalidFileObject(self):
    self.assertRaises(TypeError, yawt.YAR, 5)
    self.assertRaises(TypeError, yawt.YAR, [])

  def testInvalidFileName(self):
    invalid_filename = '/ZZZthis/is/an/invalid/filename'
    self.failIf(os.path.exists(invalid_filename))
    self.assertRaises(IOError, yawt.YAR, invalid_filename)


class YARConstructorModeTest(YARStringObjectTestCase):
  def testOpenFileForReading(self):
    yar_content = ('<YAR 1>\n'
        '<file name="test.txt">\n'
        'hello world\n'
        '</file>\n')
    yar_string = self.OpenFileString(yar_content, 'r')
    yar_object = self.OpenFileObject(yar_content, 'r')
    for yar in yar_string, yar_object:
      self.VerifyListAndRead(yar, self.file_list, self.file_contents)
      yar.Close()

  def testOpenFileForWriting(self):
    yar_content = ('<YAR 1>\n'
        '<file name="test.txt">\n'
        'hello world\n'
        '</file>\n')
    # contents of files opened in write mode will be overwritten
    yar_string = self.OpenFileString(yar_content, 'w')
    yar_object = self.OpenFileObject(yar_content, 'w')
    for yar in yar_string, yar_object:
      self.assertEqual(yar.List(), [])
      yar.Close()

  def testOpenFileForAppending(self):
    yar_content = ('<YAR 1>\n'
        '<file name="test.txt">\n'
        'hello world\n'
        '</file>\n')
    yar_string = self.OpenFileString(yar_content, 'a')
    yar_object = self.OpenFileObject(yar_content, 'a')
    for yar in yar_string, yar_object:
      self.VerifyListAndRead(yar, self.file_list, self.file_contents)
      yar.Close()

  def testInvalidMode(self):
    yar_content = ('<YAR 1>\n'
        '<file name="test.txt">\n'
        'hello world\n'
        '</file>\n')
    self.assertRaises(IOError, self.OpenFileObject, yar_content, 'foo')
    self.assertRaises(IOError, self.OpenFileString, yar_content, 'rb')


class YARConstructorEncodingTest(YARStringObjectTestCase):
  def testOpenFileEncoding(self):
    yar_content = ('<YAR 1>\n'
        '<file name="test.txt">\n'
        'hello world\n'
        '</file>\n')
    yar_string = self.OpenFileString(yar_content, 'r', binary_encoding="hex")
    yar_object = self.OpenFileObject(yar_content, 'r', binary_encoding="hex")
    for yar in yar_string, yar_object:
      self.VerifyListAndRead(yar, self.file_list, self.file_contents)
      yar.Close()

  def testInvalidEncoding(self):
    yar_content = ('<YAR 1>\n'
        '<file name="test.txt">\n'
        'hello world\n'
        '</file>\n')
    error = yawt.YARFormatError
    self.assertRaises(error, self.OpenFileObject, yar_content,
        binary_encoding='no_such_encoding')
    self.assertRaises(error, self.OpenFileObject, yar_content,
        binary_encoding='utf8')
    self.assertRaises(error, self.OpenFileObject, yar_content, mode='r',
        binary_encoding='bad')


class YARClosedFileTest(YARTestCase):
  def testListAndReadClosedFile(self):
    yar_content = ('<YAR 1>\n'
        '<file name="test.txt">\n'
        'hello world\n'
        '</file>\n')
    yar = self.OpenFileObject(yar_content)
    yar.Close()
    self.assertRaises(IOError, yar.List)
    self.assertRaises(IOError, yar.Read, 'test.txt')

  def testWriteClosedFile(self):
    yar_content = '<YAR 1>\n'
    yar = self.OpenFileObject(yar_content, 'w')
    yar.Close()
    out_file = StringIO.StringIO('test')
    self.assertRaises(IOError, yar.Write, out_file, 'test.txt')


class YARReadTest(YARTestCase):
  def testReadValidFile(self):
    yar_content = ('<YAR 1>\n'
        '<file name="eng.txt">\n'
        'quality\n'
        '</file>\n'
        '<file name="fra.txt">\n'
        'qualit\xc3\xa9\n'
        '</file>\n'
        '<file name="ita.txt">\n'
        'qualit\xc3\xa0\n'
        '</file>\n')
    yar_object = self.OpenFileObject(yar_content)
    self.assertEqual(yar_object.Read('eng.txt'), 'quality\n')
    self.assertEqual(yar_object.Read('fra.txt'), 'qualit\xc3\xa9\n')
    self.assertEqual(yar_object.Read('ita.txt'), 'qualit\xc3\xa0\n')

  def testReadInvalidFile(self):
    yar_content = ('<YAR 1>\n'
        '<file name="test.txt">\n'
        'hello world\n'
        '</file>\n')
    yar_object = self.OpenFileObject(yar_content)
    self.assertEqual(yar_object.Read('invalid.txt'), None)
    self.assertEqual(yar_object.Read('/invalid.txt'), None)
    self.assertEqual(yar_object.Read('foo/../invalid.txt'), None)


class YARMultiFilesTest(YARTestCase):
  def setUp(self):
    self.file_list = [('foo.txt', 'file'), ('bar.txt', 'file'),
      ('baz.txt', 'file'), ('all.txt', 'file')]
    self.file_contents = {'foo.txt': 'test foo\n',
      'bar.txt': 'test bar\n',
      'baz.txt': 'test baz\n',
      'all.txt': 'foo\nbar\nbaz\n\n\n'}

  def testMultiFiles(self):
    yar_content = ('<YAR 1>\n'
        '<file name="foo.txt">\n'
        'test foo\n'
        '</file>\n'
        '<file name="bar.txt">\n'
        'test bar\n'
        '</file>\n'
        '<file name="baz.txt">\n'
        'test baz\n'
        '</file>\n'
        '<file name="all.txt">\n'
        'foo\n'
        'bar\n'
        'baz\n\n\n'
        '</file>\n')
    yar = self.OpenFileObject(yar_content)
    self.VerifyListAndRead(yar, self.file_list, self.file_contents)

  def testMultiFilesWithIgnoredText(self):
    yar_content = ('<YAR 1>\n'
        'this text will get ignored\n\n'
        '<file name="foo.txt">\n'
        'test foo\n'
        '</file>\n'
        'this text will get ignored\n\n'
        '<file name="bar.txt">\n'
        'test bar\n'
        '</file>\n'
        '<file name="baz.txt">\n'
        'test baz\n'
        '</file>\n'
        '<file name="all.txt">\n'
        'foo\n'
        'bar\n'
        'baz\n\n\n'
        '</file>\n')
    yar = self.OpenFileObject(yar_content)
    self.VerifyListAndRead(yar, self.file_list, self.file_contents)

  def testMultiFilesWithDuplicateNames(self):
    yar_content = ('<YAR 1>\n'
        '<file name="duplicate.txt">\n'
        'test foo\n'
        '</file>\n'
        '<file name="duplicate.txt">\n'
        'test bar\n'
        '</file>\n')
    self.assertRaises(yawt.YARFormatError, self.OpenFileObject, yar_content)


class YARHeaderTest(YARTestCase):
  def testGoodHeader(self):
    self.assert_(self.OpenFileObject('<YAR 1>\n'))
    self.assert_(self.OpenFileObject('<YAR 1>\n<file name="foo">\n</file>'))

  def testBadHeaderFormat(self):
    error = yawt.YARFormatError
    self.assertRaises(error, self.OpenFileObject, '<YAR  1>\n')
    self.assertRaises(error, self.OpenFileObject, '<YAR    1   >\n')
    self.assertRaises(error, self.OpenFileObject, '<YAR 1 5>\n')
    self.assertRaises(error, self.OpenFileObject, '<YAR 1 5  >\n')
    self.assertRaises(error, self.OpenFileObject,
        '<YAR 1><file name="foo">\n</file>')

  def testBadHeaderVersion(self):
    error = yawt.YARFormatError
    self.assertRaises(error, self.OpenFileObject, '<YAR>\n')
    self.assertRaises(error, self.OpenFileObject, '<YAR -9>\n')
    self.assertRaises(error, self.OpenFileObject, '<YAR XYZ>\n')

  def testUnsupportedVersion(self):
    error = yawt.YARVersionError
    self.assertRaises(error, self.OpenFileObject, '<YAR 5>\n')
    self.assertRaises(error, self.OpenFileObject, '<YAR 23>\n')


class YARFileTagTest(YARTestCase):
  def setUp(self):
    self.file_list = [('test.txt', 'file')]
    self.file_contents = {'test.txt': 'hello world\n'}

  def GenerateStartTag(self, attributes):
    return ('<YAR 1>\n'
          '<file %s>\n'
          'hello world\n'
          '</file>\n' % attributes)

  def testGoodStartTags(self):
    good_start_tags = ["name='test.txt'", "name=test.txt",
        "'name'=\"test.txt\"", 'name     ="test.txt"     ',
        "'name'   =   test.txt", "'name\t'= test.txt"]
    for tag in good_start_tags:
      yar_content = self.GenerateStartTag(tag)
      yar = self.OpenFileObject(yar_content)
      self.VerifyListAndRead(yar, self.file_list, self.file_contents)

  def testBadStartTagsError(self):
    bad_start_tags = ["name\\t=test.txt", "name=", "name=   ",
        "name=test.txt this should not be here"]
    for tag in bad_start_tags:
      yar_content = self.GenerateStartTag(tag)
      self.assertRaises(yawt.YARFormatError, self.OpenFileObject, yar_content)

  def testBadStartTagsIgnored(self):
    bad_start_tags = ["name = \n test.txt", "\nname=test.txt",
        "name = \\\ntest.txt"]
    for tag in bad_start_tags:
      yar_content = self.GenerateStartTag(tag)
      yar = self.OpenFileObject(yar_content)
      self.assertEqual(yar.List(), [])

  def testBadStopTags(self):
    bad_stop_tags = ['</ file>', '< / file>', '< /FILE>', '[/FILE]',
        '</file', '</filetag>', '</"FILE">', '</fi\nle>', '</file\n>']
    for tag in bad_stop_tags:
      yar_content = ('<YAR 1>\n'
          '<file name=test.txt>\n'
          'hello world\n'
          '%s\n' % tag)
      yar = self.OpenFileObject(yar_content)
      self.assertEqual(yar.List(), [])


class YARAttributeTest(YARTestCase):
  def setUp(self):
    self.file_list = [('test.txt', 'file')]
    self.file_contents = {'test.txt': 'hello world\n'}

  def GenerateStartTag(self, attributes):
    return ('<YAR 1>\n'
          '<file %s>\n'
          'hello world\n'
          '</file>\n' % attributes)

  def testGoodAttributes(self):
    good_tags = ["name='test.txt' type=\"file\"",
        "type=file name=test.txt", "type=file name=test.txt encoding=utf8"]
    for tag in good_tags:
      yar_content = self.GenerateStartTag(tag)
      yar = self.OpenFileObject(yar_content)
      self.VerifyListAndRead(yar, self.file_list, self.file_contents)

  def testIgnoredAttributes(self):
    ignored_tags = ["name=test.txt foo=bar", "name=test.txt 5  =  5 type=file"]
    for tag in ignored_tags:
      yar_content = self.GenerateStartTag(tag)
      yar = self.OpenFileObject(yar_content)
      self.VerifyListAndRead(yar, self.file_list, self.file_contents)


class YARFileNameTest(YARTestCase):
  def GenerateStartTag(self, filename):
    return ('<YAR 1>\n'
          '<file name=%s>\n'
          'hello world\n'
          '</file>\n' % filename)

  def testGoodPaths(self):
    good_paths = {"path/to/file": "path/to/file",
        "foo/bar": "foo/bar",
        "'file name with space.txt'": "file name with space.txt",
        "'veni, vidi, vici'": "veni, vidi, vici",
        "'path/ to / file'": "path/ to / file"}
    for k in good_paths:
      v = good_paths[k]
      yar_content = self.GenerateStartTag(k)
      yar = self.OpenFileObject(yar_content)
      self.VerifyListAndRead(yar, [(v, 'file')], {v: 'hello world\n'})

  def testAbsolutePaths(self):
    absolute_paths = { "/foo": "foo", "/foo/bar/../baz": "foo/baz",
        "/////foo//////": "foo", "/foo///bar": "foo/bar",
        "'/foo bar/baz'": "foo bar/baz"}
    for k in absolute_paths:
      v = absolute_paths[k]
      yar_content = self.GenerateStartTag(k)
      yar = self.OpenFileObject(yar_content)
      self.VerifyListAndRead(yar, [(v, 'file')], {v: 'hello world\n'})

  def testRelativePaths(self):
    relative_paths = {"foo/bar/../../baz": "baz", "foo/.././././bar": "bar",
        "./foo/bar/": "foo/bar"}
    for k in relative_paths:
      v = relative_paths[k]
      yar_content = self.GenerateStartTag(k)
      yar = self.OpenFileObject(yar_content)
      self.VerifyListAndRead(yar, [(v, 'file')], {v: 'hello world\n'})

  def testBadRelativePaths(self):
    bad_relative_paths = ["../foo", "'./foo bar/../../foo'",
        "'foo/././../../bar/baz"]
    for path in bad_relative_paths:
      yar_content = self.GenerateStartTag(path)
      self.assertRaises(yawt.YARFormatError, self.OpenFileObject, yar_content)

  def testInvalidCharactersInPaths(self):
    bad_paths = ["foo/bar\\baz", "foo\\bar", "c:/foo/", "c:\\foo", "/foo/a:z"]
    for path in bad_paths:
      yar_content = self.GenerateStartTag(path)
      self.assertRaises(yawt.YARFormatError, self.OpenFileObject, yar_content)


class YAREndOfFileTest(YARTestCase):
  def testGoodEOF(self):
    good_eofs = {"FOO": "FOO",
      "'ignore me'": "ignore me",
      "'  IGNORE_ME!  '": "IGNORE_ME!"}
    for k in good_eofs:
      v = good_eofs[k]
      yar_content = ('<YAR 1>\n'
          '<file name=test.txt eof=%s>\n'
          'file includes the next two lines\n'
          '</file>\n'
          '</file IGNORE_ME>\n'
          '</file %s>\n' % (k, v))
      yar = self.OpenFileObject(yar_content)
      self.VerifyListAndRead(yar, [('test.txt', 'file')],
          {'test.txt': ('file includes the next two lines\n'
          '</file>\n'
          '</file IGNORE_ME>\n')})

  def testBadEOF(self):
    yar_content = ('<YAR 1>\n'
        '<file name=test.txt eof=FOO>\n'
        'this file only contain this line\n'
        '</file FOO>\n'
        'the rest of the file should be ignored\n'
        '</file>\n')
    yar = self.OpenFileObject(yar_content)
    self.VerifyListAndRead(yar, [('test.txt', 'file')],
        {'test.txt': 'this file only contain this line\n'})


class YARBinaryFilesTest(YARTestCase):
  def testBase64Encoding(self):
    yar_content = ('<YAR 1>\n'
        '<file name=base64.txt encoding=base64>\n'
        'dGVzdAoxMjMK\n'
        '</file>\n'
        'the rest of the file should be ignored\n'
        '</file>\n')
    yar = self.OpenFileObject(yar_content)
    self.VerifyListAndRead(yar, [('base64.txt', 'file')],
        {'base64.txt': 'test\n123\n'})

  def testHexEncoding(self):
    yar_content = ('<YAR 1>\n'
        '<file name=hex.txt encoding=hex>\n'
        '746573740a3132330a\n'
        '</file>\n')
    yar = self.OpenFileObject(yar_content)
    self.VerifyListAndRead(yar, [('hex.txt', 'file')],
        {'hex.txt': 'test\n123\n'})

  def testInvalidEncoding(self):
    yar_content = ('<YAR 1>\n'
        '<file name=invalid.txt encoding=invalidFoo>\n'
        '746573740a3132330a\n'
        '</file>\n')
    self.assertRaises(yawt.YARFormatError, self.OpenFileObject, yar_content)

  def testDosFiles(self):
    file_list = [('test_base64.bin', 'file'), ('test_hex.bin', 'file')]
    binary_content = '\x00\x10\x20\r\x30\x40\x50\r\n\xff\xfe\n\xfd'
    file_contents = {'test_base64.bin': binary_content,
      'test_hex.bin': binary_content}
    yar_content = ('<YAR 1>\r\n'
        '<file name="test_base64.bin" encoding="base64">\r\n'
        'ABAgDTBAUA0K//4K/Q==\r\n'
        '</file>\r\n'
        '<file name="test_hex.bin" encoding="hex">\r\n'
        '001020\r\n'
        '0d3040500d0afffe0afd\r\n'
        '</file>\r\n')
    yar_object = self.OpenFileObject(yar_content)
    self.VerifyListAndRead(yar_object, file_list, file_contents)


class YARFileTypeTest(YARTestCase):
  def testDirectory(self):
    yar_content = ('<YAR 1>\n'
        '<file name=testdir type=directory>\n'
        '</file>\n')
    yar = self.OpenFileObject(yar_content)
    self.VerifyListAndRead(yar, [('testdir', 'directory')], {'testdir': ''})

  def testDirectoryWithContent(self):
    yar_content = ('<YAR 1>\n'
        '<file name=testdir type=directory>\n'
        'this gets ignored\n'
        '</file>\n')
    yar = self.OpenFileObject(yar_content)
    self.VerifyListAndRead(yar, [('testdir', 'directory')], {'testdir': ''})

  def testDirectoryWithAttributes(self):
    yar_content = ('<YAR 1>\n'
        '<file name=testdir type=directory encoding=hex eof=FOO>\n'
        '</file FOO>\n')
    yar = self.OpenFileObject(yar_content)
    self.VerifyListAndRead(yar, [('testdir', 'directory')], {'testdir': ''})

  def testFile(self):
    yar_content = ('<YAR 1>\n'
        '<file name=test.txt type=file>\n'
        'this does not gets ignored\n'
        '</file>\n')
    yar = self.OpenFileObject(yar_content)
    self.VerifyListAndRead(yar, [('test.txt', 'file')],
        {'test.txt': 'this does not gets ignored\n'})

  def testInvalid(self):
    yar_content = ('<YAR 1>\n'
        '<file name=invalid.txt type=invalidFoo>\n'
        'type is invalid\n'
        '</file>\n')
    self.assertRaises(yawt.YARFormatError, self.OpenFileObject, yar_content)


class YARWriteTest(YARTestCase):
  def setUp(self):
    self.test_file1_content = 'test file 1\n'
    self.test_file1 = StringIO.StringIO(self.test_file1_content)
    self.test_file2_content = '\101\102\103\5\0'
    self.test_file2 = StringIO.StringIO(self.test_file2_content)
    self.yar_content = ('<YAR 1>\n'
        '<file name=testdir type=directory>\n'
        '</file>\n')

  def testWriteFile(self):
    yar = self.OpenFileObject(self.yar_content, 'w')
    self.VerifyListAndRead(yar, [], {})
    yar.Write(self.test_file1, 'testfile')
    self.VerifyListAndRead(yar, [('testfile', 'file')],
        {'testfile': self.test_file1_content})

  def testAppendFile(self):
    yar_list = [('testdir', 'directory'), ('testfile', 'file')]
    yar_files = {'testdir': '', 'testfile': self.test_file1_content}

    yar = self.OpenFileObject(self.yar_content, 'a')
    self.VerifyListAndRead(yar, [('testdir', 'directory')], {'testdir': ''})
    yar.Write(self.test_file1, 'testfile')
    self.VerifyListAndRead(yar, yar_list, yar_files)

  def testWriteToReadOnlyFile(self):
    yar = self.OpenFileObject(self.yar_content, 'r')
    self.assertRaises(IOError, yar.Write, self.test_file1,
        'testfile')

  def testWriteInvalidFilename(self):
    yar = self.OpenFileObject(self.yar_content, 'a')
    self.assertRaises(yawt.YARFormatError, yar.Write, self.test_file1,
        'c:\\badname')

  def testWriteDuplicateFilename(self):
    yar = self.OpenFileObject(self.yar_content, 'a')
    self.assertRaises(yawt.YARFormatError, yar.Write, self.test_file1,
        'testdir')

  def testWriteInvalidOutfile(self):
    yar = self.OpenFileObject(self.yar_content, 'a')
    self.assertRaises(TypeError, yar.Write, None, 'newfile')
    self.assertRaises(TypeError, yar.Write, [], 'newfile')

  def testWriteWithDefaultEncoding(self):
    yar_list = [('testdir', 'directory'), ('testfile.bin', 'file')]
    yar_files = {'testdir': '', 'testfile.bin': self.test_file2_content}

    yar = self.OpenFileObject(self.yar_content, 'a')
    yar.Write(self.test_file2, 'testfile.bin')
    self.VerifyListAndRead(yar, yar_list, yar_files)

  def testWriteWithValidEncoding(self):
    yar_list = [('testdir', 'directory'), ('testfile.bin', 'file')]
    yar_files = {'testdir': '', 'testfile.bin': self.test_file2_content}

    yar = self.OpenFileObject(self.yar_content, 'a')
    yar.Write(self.test_file2, 'testfile.bin', 'hex')
    self.VerifyListAndRead(yar, yar_list, yar_files)

  def testWriteWithInvalidEncoding(self):
    yar = self.OpenFileObject(self.yar_content, 'a')
    yar.Write(self.test_file2, 'testfile.bin', 'hex')
    self.assertRaises(yawt.YARFormatError, yar.Write, self.test_file2,
        'testfile.bin', 'bad_encoding')

  def testWriteDirectory(self):
    yar_list = [('testdir', 'directory')]
    yar_files = {'testdir': ''}

    yar = self.OpenFileObject(self.yar_content, 'a')
    self.VerifyListAndRead(yar, yar_list, yar_files)
    yar.Write(None, 'newdir1', is_dir=True)
    yar_list.append(('newdir1', 'directory'))
    yar_files['newdir1'] = ''
    self.VerifyListAndRead(yar, yar_list, yar_files)
    yar.Write(self.test_file2, 'newdir2', is_dir=True)
    yar_list.append(('newdir2', 'directory'))
    yar_files['newdir2'] = ''
    self.VerifyListAndRead(yar, yar_list, yar_files)
    yar.Write(None, 'newdir3', binary_encoding='hex', is_dir=True)
    yar_list.append(('newdir3', 'directory'))
    yar_files['newdir3'] = ''
    self.VerifyListAndRead(yar, yar_list, yar_files)

  def testWriteDosFile(self):
    content = '<file>\ntest 1 2 3\n</file>\n'
    file_with_cr = StringIO.StringIO(content.replace('\n', '\r\n'))
    yar_list = [('testdir', 'directory'), ('testfile.dos', 'file')]
    yar_files = {'testdir': '', 'testfile.dos': content}
    yar = self.OpenFileObject(self.yar_content, 'a')
    yar.Write(file_with_cr, 'testfile.dos')
    self.VerifyListAndRead(yar, yar_list, yar_files)


class IsBinaryTest(unittest.TestCase):
  def testBinaryData(self):
    binary_data = ['\0', 'test\255', '\xf1 != n?',
      '\x7f\x45\x4c\x46\x01\x01\x01\x00']
    for d in binary_data:
      self.assertEqual(yawt.IsBinary(d), True)

  def testTextData(self):
    text_data = ['test', 'foo', 'ABC__123~!@abc', '!@#$%^&*()', '"', '']
    for d in text_data:
      self.assertEqual(yawt.IsBinary(d), False)

  def testUnicodeTextData(self):
    unicode_text_data = ['\xe6\xa2\x81\xe4\xb9\x83\xe5\xbc\xb7']
    for d in unicode_text_data:
      self.assertEqual(yawt.IsBinary(d), False)

  def testInvalidData(self):
    invalid_data = [None, {}, ['data']]
    for d in invalid_data:
      self.assertRaises(TypeError, yawt.IsBinary, d)


class IsFileLikeObjectTest(unittest.TestCase):
  def setUp(self):
    (fd, self.tmpname) = tempfile.mkstemp('yartest')
    os.close(fd)
    self.tmpfile = open(self.tmpname)

  def tearDown(self):
    self.tmpfile.close()
    os.remove(self.tmpname)

  def testFileLikeObjects(self):
    file_objs = [self.tmpfile, StringIO.StringIO('foo')]
    for f in file_objs:
      self.assertEqual(yawt.IsFileLikeObject(f), True)

  def testNonFileLikeObjects(self):
    bad_objs = [None, "", {}]
    for f in bad_objs:
      self.assertEqual(yawt.IsFileLikeObject(f), False)


class YAWTInstalldataTestCase(unittest.TestCase):
  """
  section -> windows "key"
  key -> windows "value"
  value -> windows "data"

  sections cannot have '\' in name, '\' only used as separator
  all keys should be in double quotes, except the 'default key', which is just
    @ without any quotes

  REG_BINARY starts with hex:
  REG_DWORD starts with dword:
  REG_MULTI_SZ starts with hex(7):
  REG_EXPAND_SZ starts with hex(2):
  REG_NONE starts with hex(0):
  REG_SZ data is in double quotes

  comment lines start with ; or #
  \ indicates line wraps
  """
  def ValidateSections(self, parser, correct_data):
    self.assertEqual(parser.keys(), correct_data.keys())
    for s in correct_data.keys():
      self.assertEqual(parser[s], correct_data[s])


class YAWTInstalldataSectionsTest(YAWTInstalldataTestCase):
  def testValidSections(self):
    sections = {'Test\\Section':
      {None: ('Default Value', 'REG_SZ'),
        'Key1': ('\x10\x12\x23', 'REG_BINARY')},
      'Yet\\Another\\\'"[]=/@Section':
        {'FOO\'"\\[]=/': ('BAR"\\[]=/', 'REG_SZ'), '@': ('Foo', 'REG_SZ')}}

    data_sections_with_separation = ('[Test\\Section]\n'
        '@="Default Value"\n'
        '"Key1"=hex:10,12,23\n'
        '\n'
        '\n'
        '[Yet\\Another\\\'"[]=/@Section]\n'
        '"FOO\'\\"\\\\[]=/"="BAR\\"\\\\[]=/"\n'
        '"@"="Foo"\n')
    data_sections_without_separation = ('[Test\\Section]\n'
        '@="Default Value"\n'
        '"Key1"=hex:10,12,23\n'
        '[Yet\\Another\\\'"[]=/@Section]\n'
        '"FOO\'\\"\\\\[]=/"="BAR\\"\\\\[]=/"\n'
        '"@"="Foo"\n')
    for f in [data_sections_without_separation, data_sections_with_separation]:
      w = yawt.WinRegParser(f)
      self.ValidateSections(w, sections)

  def testEmptySections(self):
    sections = {'section 1': {}, 'section 2': {}, 'section 3': {}}
    data_empty_sections = ('[section 1]\n'
        '[section 2]\n'
        '\n'
        '[section 3]\n')
    w = yawt.WinRegParser(data_empty_sections)
    self.ValidateSections(w, sections)

  def testBadlyFormattedSections(self):
    data_missing_left_bracket = ('Test\\Section]\n'
        '@="Default Value"\n'
        '"Test Key"="Foo"\n')
    data_misplaced_right_bracket = ('[Test\\Section\n'
        ']@="Default Value"\n'
        '"Test Key"="Foo"\n')
    data_missing_newline = ('[Test\\Section]'
        '@="Default Value"\n'
        '"Test Key"="Foo"\n')
    data_two_consecutive_backslashes = ('[Test\\\\Section]\n'
        '@="Default Value"\n'
        '"Test Key"="Foo"\n')
    data_begin_with_backslash = ('[\\Test\\Section]\n'
        '@="Default Value"\n'
        '"Test Key"="Foo"\n')
    data_end_with_backslash = ('[Test\\Section\\]\n'
        '@="Default Value"\n'
        '"Test Key"="Foo"\n')
    data_no_section_name = ('[]\n'
        '@="Default Value"\n'
        '"Test Key"="Foo"\n')
    for f in [data_missing_left_bracket, data_misplaced_right_bracket,
      data_missing_newline, data_two_consecutive_backslashes,
      data_begin_with_backslash, data_end_with_backslash, data_no_section_name]:
      self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)

  def testNoData(self):
    for f in [None, '', '\n\n\n']:
      w = yawt.WinRegParser(f)
      self.assertEqual(w.keys(), [])

class YAWTInstalldataKeysTest(YAWTInstalldataTestCase):
  def testDuplicateKeys(self):
    data_duplicate_key = ('[Test\\Section]\n'
        '"Duplicate Key"="Foo"\n'
        '"Duplicate Key"="Bar"\n')
    data_duplicate_default_value = ('[Test\\Section]\n'
        '@="Foo"\n'
        '@="Foo"\n'
        '"Test Key"="Foo"\n')
    for f in [data_duplicate_key, data_duplicate_default_value]:
      self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)

  def testInvalidKeys(self):
    data_empty_key = ('[Test Section]\n'
        '@="Foo"\n'
        '""="Bar"\n')
    data_invalid_key1 = ('[Test Section]\n'
        '@="Foo"\n'
        '500="Bar"\n')
    data_invalid_key2 = ('[Test Section]\n'
        '@="Foo"\n'
        '[]="Bar"\n')
    for f in [data_empty_key, data_invalid_key1, data_invalid_key2]:
      self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)

  def testKeyWithoutSection(self):
    data_no_section = ('"Key1"="Foo"\n'
        '"Key2"="Bar"\n')
    data_comment_instead_of_section = ('#[no section here]\n'
        '@="Foo"\n'
        '"Test Key"="Foo"\n')
    for f in [data_no_section, data_comment_instead_of_section]:
      self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)


class YAWTInstalldataQuotesTest(YAWTInstalldataTestCase):
  def testValidUnquotedValues(self):
    sections = {'Test\\Section':
      {'Binary': ('\x10\x12\x23', 'REG_BINARY'),
      'Multi_SZ': ('\x10\x10\x10\x00\x15\x15\x15\x00\x00', 'REG_MULTI_SZ'),
      'Expand_SZ': ('\x30\x30\x30\x00', 'REG_EXPAND_SZ'),
      'Dword': (0xffffffff, 'REG_DWORD'),
      'SZ': ('strings should be quoted', 'REG_SZ')}}
    data_valid_unquoted = ('[Test\\Section]\n'
        '"Binary"=hex:10,12,23\n'
        '"Multi_SZ"=hex(7):10,10,10,00,15,15,15,00,00\n'
        '"Expand_SZ"=hex(2):30,30,30,00\n'
        '"Dword"=dword:ffffffff\n'
        '"SZ"="strings should be quoted"\n')
    w = yawt.WinRegParser(data_valid_unquoted)
    self.ValidateSections(w, sections)

  def testInvalidQuotes(self):
    data_unquoted_key = ('[Test\\Section]\n'
        '@="Default Value"\n'
        'Key="Foo"\n')
    data_unquoted_value = ('[Test\\Section]\n'
        '@="Default Value"\n'
        '"Test Key"=Foo\n')
    data_missing_left_quote = ('[Test\\Section]\n'
        '@="Default Value"\n'
        'Test Key"="Foo"\n')
    data_missing_right_quote = ('[Test\\Section]\n'
        '@="Default Value"\n'
        '"Test Key="Foo"\n')
    data_unescaped_quote = ('[Test\\Section]\n'
        '@="Default Value"\n'
        '"Test"Key="Foo"\n')
    for f in [data_unquoted_key, data_unquoted_value, data_missing_left_quote,
      data_missing_right_quote, data_unescaped_quote]:
      self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)


class YAWTInstalldataValueTest(YAWTInstalldataTestCase):
  def testInvalidValues(self):
    data_invalid_value1 = ('[Test Section]\n'
        'data=xyzzy:1000')
    data_invalid_value2 = ('[Test Section]\n'
        'data=hex(3):10,10,00')
    data_invalid_value3 = ('[Test Section]\n'
        'data=hex(900):10,10,00')
    data_invalid_value4 = ('[Test Section]\n'
        'data=hex(xyzzy):10,10,00')
    for f in [data_invalid_value1, data_invalid_value2, data_invalid_value3,
      data_invalid_value4]:
        self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)

  def testInvalidDwordValues(self):
    data_invalid_dword1 = ('[Test\\Section]\n'
        '"too long"=dword:1111111111111111\n')
    data_invalid_dword2 = ('[Test\\Section]\n'
        '"invalid char"=dword:123x\n')
    data_invalid_dword3 = ('[Test\\Section]\n'
        '"no value"=dword:\n')
    for f in [data_invalid_dword1, data_invalid_dword2, data_invalid_dword3]:
      self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)

  def testValidBinaryValues(self):
    sections = {'Test Section':
      {'no value': ('', 'REG_BINARY'), 'binary1': ('\x01', 'REG_BINARY')}}
    data_binary = ('[Test Section]\n'
        '"no value"=hex:\n'
        '"binary1"=hex:01\n')
    w = yawt.WinRegParser(data_binary)
    self.ValidateSections(w, sections)

  def testInvalidBinaryValues(self):
    data_invalid_binary1 = ('[Test\\Section]\n'
        '"single digit"=hex:10,10,2,10\n')
    data_invalid_binary2 = ('[Test\\Section]\n'
        '"missing comma"=hex:10,10,1010\n')
    data_invalid_binary3 = ('[Test\\Section]\n'
        '"invalid hex"=hex:10,10,gg,10\n')
    for f in [data_invalid_binary1, data_invalid_binary2, data_invalid_binary3]:
        self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)

  def testValidNoneValues(self):
    sections = {'Test Section':
      {'no value': ('', 'REG_NONE'), 'none1': ('\x01', 'REG_NONE'),
      'none3': ('\x11\x12\x13', 'REG_NONE')}}
    data_binary = ('[Test Section]\n'
        '"no value"=hex(0):\n'
        '"none1"=hex(0): 01\n'
        '"none3"=hex(0):11, 12, 13\n')
    w = yawt.WinRegParser(data_binary)
    self.ValidateSections(w, sections)

  def testInvalidNoneValues(self):
    data_invalid_none1 = ('[Test\\Section]\n'
        '"single digit"=hex:10,10,10,2\n')
    data_invalid_none2 = ('[Test\\Section]\n'
        '"missing comma"=hex:10,1010,10\n')
    data_invalid_none3 = ('[Test\\Section]\n'
        '"invalid hex"=hex:10,10,10,xy\n')
    for f in [data_invalid_none1, data_invalid_none2, data_invalid_none3]:
        self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)

  def testValidMultiszValues(self):
    sections = {'Test Section':
      {'no characters': ('\x00', 'REG_MULTI_SZ'),
      'one character': ('\x31\00', 'REG_MULTI_SZ')}}
    data_multi = ('[Test Section]\n'
        '"no characters"=hex(7):00\n'
        '"one character"=hex(7):31,00\n')
    w = yawt.WinRegParser(data_multi)
    self.ValidateSections(w, sections)

  def testInvalidMultiszValues(self):
    data_invalid_multisz1 = ('[Test\\Section]\n'
        '"single digit"=hex(7):10,10,10,0,20\n')
    data_invalid_multisz2 = ('[Test\\Section]\n'
        '"no value "=hex(7):\n')
    data_invalid_multisz3 = ('[Test\\Section]\n'
        '"trailing comma"=hex(7):10,10,10,00,\n')
    data_invalid_multisz4 = ('[Test\\Section]\n'
        '"invalid hex"=hex(7):1010,ggg\n')
    for f in [data_invalid_multisz1, data_invalid_multisz2,
      data_invalid_multisz3, data_invalid_multisz4]:
        self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)

  def testValidExpandszValues(self):
    sections = {'Test Section':
      {'no characters': ('\x00', 'REG_EXPAND_SZ'),
      'one character': ('\x31\x00', 'REG_EXPAND_SZ')}}
    data_expand = ('[Test Section]\n'
        '"no characters"=hex(2):00\n'
        '"one character"=hex(2):31,00\n')
    w = yawt.WinRegParser(data_expand)
    self.ValidateSections(w, sections)

  def testInvalidExpandszValues(self):
    data_invalid_expandsz1 = ('[Test\\Section]\n'
        '"single digit"=hex(2):05,5,00\n')
    data_invalid_expandsz2 = ('[Test\\Section]\n'
        '"no value "=hex(2):\n')
    data_invalid_expandsz3 = ('[Test\\Section]\n'
        '"trailing comma"=hex(2):30,30,00,\n')
    data_invalid_expandsz4 = ('[Test\\Section]\n'
        '"invalid hex"=hex(2):10,20,5!,00\n')
    for f in [data_invalid_expandsz1, data_invalid_expandsz2,
      data_invalid_expandsz3, data_invalid_expandsz4]:
        self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)


class YAWTInstalldataCommentsTest(YAWTInstalldataTestCase):
  def testValidComments(self):
    sections = {'Test\\Section':
      {'Key1': ('Foo', 'REG_SZ'), 'Key2': ('Bar', 'REG_SZ')}}
    data_comments = ('#[this file is riddled with comments]\n'
        '; this comment starts with a semicolon\n'
        '[Test\\Section]\n'
        '"Key1"="Foo"\n'
        '; yet another comment\n'
        '"Key2"="Bar"\n')
    w = yawt.WinRegParser(data_comments)
    self.ValidateSections(w, sections)

  def testInvalidComments(self):
    data_invalid_comments = ('[Test\\Section]\n'
        '"Key1"="Foo" # this comment should not be here\n'
        '"Key2"="Bar"\n')
    self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser,
        data_invalid_comments)


class YAWTInstalldataMultiLineTest(YAWTInstalldataTestCase):
  def testValidMultiLine(self):
    sections = {'Test\\Section':
      {'linewrap': ('\x10\x12\x23\x14\x14\x11\x51\x51\x51\x51\x12\x23\x44\x55',
          'REG_BINARY'), 'nowrap': (0xffffffff, 'REG_DWORD')}}
    data_valid_multiline = ('[Test\\Section]\n'
        '"linewrap"=hex:10,12,23,14,14,11,51,51,51,51,\\\n'
        '  12,23\\\n'
        ',44,55\n'
        '"nowrap"=dword:ffffffff')
    w = yawt.WinRegParser(data_valid_multiline)
    self.ValidateSections(w, sections)

  def testInvalidMultiLine(self):
    data_no_backslash = ('[Test\\Section]\n'
        '"linewrap"=hex:10,12,23,14,14,11,51,51,51,51,\n'
        '  12,23\\\n'
        ',44,55\n'
        '"nowrap"=dword:ffffffff')
    data_extra_backslash = ('[Test\\Section]\n'
        '"linewrap"=hex:10,12,23,14,14,11, \\ 51,51,51,51,\\\n'
        '  12,23\\\n'
        ',45,55\n'
        '"nowrap"=dword:ffffffff')
    for f in [data_no_backslash, data_extra_backslash]:
      self.assertRaises(yawt.WinRegParserFormatError, yawt.WinRegParser, f)


class YAWTInstalldataWriteTest(unittest.TestCase):
  def ValidateOutput(self, reg_content):
    parser = yawt.WinRegParser(reg_content)
    self.assertEqual(parser.WriteOutput(), reg_content)

  def testNoInput(self):
    for f in [None, '', '\n\n\n']:
      parser = yawt.WinRegParser(f)
      self.assertEqual(parser.WriteOutput(), '')

  def testMultipleSections(self):
    reg = ('[Test\\Section]\n'
        '@="Default Value"\n'
        '"Key1"=hex:10,12,23\n'
        '\n'
        '[Yet\\Another\\\'"[]=/@Section]\n'
        '"@"="Foo"\n'
        '"FOO\'\\"\\\\[]=/"="BAR\\"\\\\[]=/"\n\n')
    self.ValidateOutput(reg)

  def testDataTypes(self):
    reg = ('[Test Section]\n'
        '@="Default Value"\n'
        '"binary"=hex:10,12,23,45,56\n'
        '"binary_empty"=hex:\n'
        '"binary_one"=hex:42\n'
        '"dword"=dword:deadbeef\n'
        '"expand"=hex(2):10,20,30,1a,2a,3a,1f,2f,3f,10,20,30,10,20,30,00\n'
        '"expand_empty"=hex(2):00\n'
        '"multi"=hex(7):10,12,00,20,30,00,ff,ef,a1,00,00\n'
        '"multi_empty"=hex(7):00\n'
        '"string"="test string"\n'
        '"string_empty"=""\n\n')
    self.ValidateOutput(reg)


class YAWTInstalldataAddEntryTest(YAWTInstalldataTestCase):
  def testAddValidEntries(self):
    sections = {'Test Section':
      {'Key1': ('Foo', 'REG_SZ'), 'Key 2': (0x3456, 'REG_DWORD'),
        'Key3': ('', 'REG_SZ'),
      'Key #4': (0, 'REG_DWORD'), 'Key 5': (0xffffffff, 'REG_DWORD')},
      'Another Section':
        {None: ('\x30\x31\x32\00', 'REG_EXPAND_SZ'),
        '(Default)': ('\x10\x20\x30', 'REG_BINARY'),
        'Default': ('\x30\x31\x32\00\x33\x34\x35\00\00', 'REG_MULTI_SZ'),
        'DEFAULT': ('\x00', 'REG_MULTI_SZ'),
        '@': ('\x20\x00\x00', 'REG_MULTI_SZ'),
        'none1': ('', 'REG_NONE'),
        'none2': ('\x30\x4f\xff', 'REG_NONE'),
        'binary': ('', 'REG_BINARY')}}
    parser = yawt.WinRegParser()
    parser.AddEntry('Test Section', 'Key1', 'Foo', 'REG_SZ')
    parser.AddEntry('Test Section', 'Key3', '', 'REG_SZ')
    parser.AddEntry('Test Section', 'Key 2', 0x3456, 'REG_DWORD')
    parser.AddEntry('Test Section', 'Key #4', 0, 'REG_DWORD')
    parser.AddEntry('Test Section', 'Key 5', 0xffffffff, 'REG_DWORD')
    parser.AddEntry('Another Section', None, '\x30\x31\x32', 'REG_EXPAND_SZ')
    parser.AddEntry('Another Section', '(Default)', '\x10\x20\x30',
        'REG_BINARY')
    parser.AddEntry('Another Section', 'binary', '', 'REG_BINARY')
    parser.AddEntry('Another Section', 'Default',
        ['\x30\x31\x32', '\x33\x34\x35'], 'REG_MULTI_SZ')
    parser.AddEntry('Another Section', '@', [' '], 'REG_MULTI_SZ')
    parser.AddEntry('Another Section', 'DEFAULT', [], 'REG_MULTI_SZ')
    parser.AddEntry('Another Section', 'none1', '', 'REG_NONE')
    parser.AddEntry('Another Section', 'none2', '\x30\x4f\xff', 'REG_NONE')
    self.ValidateSections(parser, sections)

  def testAddInvalidKey(self):
    parser = yawt.WinRegParser()
    for k in ['', [], 55]:
      self.assertRaises(yawt.WinRegParserFormatError, parser.AddEntry, 'test',
          k, 0x3456, 'REG_DWORD')

  def testAddInvalidSection(self):
    parser = yawt.WinRegParser()
    for s in ['', [], 55, None, '\\test', 'test\\\\section']:
      self.assertRaises(yawt.WinRegParserFormatError, parser.AddEntry, s,
          'key', 0x3456, 'REG_DWORD')

  def testAddInvalidDwordValues(self):
    parser = yawt.WinRegParser()
    for v in [None, [], -1, 0x100000000, '', 'test']:
      self.assertRaises(yawt.WinRegParserFormatError, parser.AddEntry, 'test',
          'key', v, 'REG_DWORD')

  def testAddInvalidMultiszValues(self):
    parser = yawt.WinRegParser()
    for v in [None, {}, 500, "test", 'test\x00test2']:
      self.assertRaises(yawt.WinRegParserFormatError, parser.AddEntry, 'test',
          'key', v, 'REG_MULTI_SZ')

  def testAddInvalidValues(self):
    parser = yawt.WinRegParser()
    for v in [None, 500, []]:
      self.assertRaises(yawt.WinRegParserFormatError, parser.AddEntry, 'test',
          'key', v, 'REG_BINARY')
      self.assertRaises(yawt.WinRegParserFormatError, parser.AddEntry, 'test',
          'key', v, 'REG_EXPAND_SZ')
      self.assertRaises(yawt.WinRegParserFormatError, parser.AddEntry, 'test',
          'key', v, 'REG_SZ')
      self.assertRaises(yawt.WinRegParserFormatError, parser.AddEntry, 'test',
          'key', v, 'REG_NONE')

  def testAddInvalidType(self):
    parser = yawt.WinRegParser()
    for t in [None, {}, '', 'Wrong', 'sZ', 'binarys']:
      self.assertRaises(yawt.WinRegParserFormatError, parser.AddEntry, 'test',
          'key', 'xyz', t)

  def testOverwriteExistingEntry(self):
    reg = ('[Test Section]\n'
        '@="Default Value"\n'
        '"Key1"=hex:10,12,23\n')
    sections = {'Test Section':
      {None: ('Test', 'REG_SZ'), 'Key1': (500, 'REG_DWORD')}}
    parser = yawt.WinRegParser(reg)
    parser.AddEntry('Test Section', None, 'Test', 'REG_SZ')
    parser.AddEntry('Test Section', 'Key1', 500, 'REG_DWORD')
    self.ValidateSections(parser, sections)


if __name__ == '__main__':
  unittest.main()
